#!/usr/bin/python

# hostbuilder - Tools for DNS zonefile administration
# Copyright (c) 2010 Crown copyright
# 
# This file is part of hostbuilder.
# 
# hostbuilder is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# hostbuilder is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with hostbuilder.  If not, see <http://www.gnu.org/licenses/>.


import sys
import os.path
import re

import yaml

zone_re = re.compile(r'^(?P<host>[\w.-]*)\s+IN\s+(?P<rr>[A-Z]+)\s+(?P<value>.*?)\s*$')

import re 

def naturalkey(s):
    """Return a "natural" key for use in natural sorting comparisons.""" 
    k = []
    for c in re.split(r'(\d+)', s):
        if c.isdigit():
            k.append(int(c))
        else:
            k.append(c)
    return tuple(k)


class InvalidFQDNError(Exception):
    """An incorrect or invalid FQDN"""

class InvalidZoneError(Exception):
    """The zone is invalid."""


class Zone(object):
    """Utility class representing a DNS zone file."""

    def __init__(self, name, soa=None, admin=None):
        assert bool(name)
        assert not name.endswith('.')

        self.name = name
        if admin:
            self.admin = admin.replace('@', '.')
        else:
            self.admin = 'hostmaster.' + self.name

        self.soa = soa
        self.rrs = []

    def add_from_template(self, fname):
        for l in open(fname):
            mo = zone_re.match(l)
            if mo:
                self.rrs.append((mo.group('rr'), mo.group('host'), mo.group('value')))

    def qualify(self, short_name):
        if short_name.endswith('.'):
            return short_name[:-1]
        else:
            return short_name + '.' + self.name
        
    def lookup(self, request, reqtype='A'):
        try:
            req = self.short_name(request)
        except InvalidFQDNError:
            return []

        results = []
        for type, host, value in self.rrs:
            if host in ['@', '']:
                host = '@'
            if type == reqtype and host == req:
                results.append(self.qualify(value))
        return results
    
    def get_soa(self):
        if self.soa is None:
            nameservers = self.lookup(self.name, 'NS')
            if nameservers:
                return nameservers[0]
            else:
                return self.name
        return self.soa


    def write_zonefile(self, out, serial):
        self.sort()
        out.write("""; This zone file was automatically generated with hostbuilder.py
; See http://code.google.com/p/hostbuilder/
""")
        out.write("$ORIGIN %s.\n" % self.name)
        out.write("$TTL 86400\n")
        out.write("""@                           IN    SOA       %s. %s. (
    %s  ; serial
    21600      ; refresh after 6 hours
    3600       ; retry after 1 hour
    604800     ; expire after 1 week
    86400      ; minimum TTL of 1 day 
)\n""" % (self.get_soa(), self.admin, serial))
        lasttype = None
        for type, host, value in self.rrs:
            if type != lasttype:
                out.write('\n')
                lasttype = type
            out.write("%s    IN    %s    %s\n" % (host.ljust(24), type.ljust(10), value))

        out.write('\n')

    def write_template(self, out):
        self.write_zonefile(out, '${SERIAL}')

    def sort(self):
        order = {
            'NS': 0,
            'MX': 1,
            'A': 2,
            'CNAME': 3
        }
        self.rrs.sort(key=lambda x: (order.get(x[0], x[0]), naturalkey(x[2]), naturalkey(x[1])))

    def __contains__(self, fqdn):
        return fqdn.endswith('.' + self.name)

    def short_name(self, fqdn):
        """Return a short name for a host that is expected to be within the zone."""
        if fqdn == self.name:
            return '@'
        if not fqdn.endswith('.' + self.name):
            raise InvalidFQDNError("fqdn %s not in zone %s" % (fqdn, self.name))
        return fqdn[:-len(self.name)-1]

    def qualified_name(self, fqdn):
        """Return a name for any host use in a zone file - either a short name if 
        one exists or an FQDN terminated with a dot otherwise.
        """
        try:
            return self.short_name(fqdn)
        except InvalidFQDNError:
            return fqdn + '.'

    def add_rr(self, type, fqdn, value):
        host = self.short_name(fqdn)
        self.rrs.append((type, host, value))

    def add_ns(self, nameserver_fqdn):
        self.add_rr('NS', self.name, self.qualified_name(nameserver_fqdn))

    def add_mx(self, mailserver_fqdn, priority=10):
        self.add_rr('MX', self.name, '%d %s' % (priority, self.qualified_name(mailserver_fqdn)))

    def add_a(self, fqdn, ipv4):
        self.add_rr('A', fqdn, ipv4)

    def add_cname(self, alias_fqdn, canonical_fqdn):
        self.add_rr('CNAME', alias_fqdn, self.qualified_name(canonical_fqdn))

    def check(self):
        """Check that there are no duplicate definitions of an incompatible type"""
        records = {}
        for type, host, value in self.rrs:
            if host in records:
                otype = records[host]
                if 'CNAME' in [type, otype]:
                    raise InvalidZoneError("CNAME and some other data for host %s" % host)
            records[host] = type 


class ReverseZone(Zone):
    def sort(self):
        def _key(rr):
            try:
                k = [int(c) for c in rr[1].split('.')]
                k.reverse()
            except ValueError:
                k = [0, rr[1]]
            return k
            
        self.rrs.sort(key=_key)

    def add_ptr(self, reverse_ip, fqdn):
        self.add_rr('PTR', reverse_ip, fqdn + '.')


class Host(object):
    def __init__(self, canonical, ips, aliases, public=None, comment=None):
        self.canonical = canonical
        self.ips = ips
        self.aliases = aliases
        self.public = public
        self.comment = comment

    def reverse_ips(self):
        rips = []
        for ip in self.ips:
           parts = ip.split('.')
           parts.reverse()
           rip = '.'.join(parts) + '.in-addr.arpa'
           rips.append(rip)
        return rips

    def routeable_ip(self, routes):
        """Return the first IP for which any route is available"""
        for i in self.ips:
            for r in routes:
                if i in r:
                    return i
        raise NoRouteToHost(self.canonical)


class NoRouteToHost(Exception):
    """No IP is routeable"""

class HostSyntaxError(Exception):
    """A problem loading the YAML zonefile"""


def load_hosts(hostfile='hosts.yaml'):
    hosts = yaml.load(open(hostfile))

    KNOWN_KEYS = set(['ips', 'aliases', 'comment', 'public'])

    hs = []
    for canonical, props in hosts.items():
        for k in props:
            if k not in KNOWN_KEYS:
                raise HostSyntaxError("Unknown property %s in host %s." % (k, canonical))
        try:
            ips = props['ips']
        except KeyError:
            raise HostSyntaxError("No IPs listed for host %s." % canonical)

        if isinstance(ips, str):
            ips = [ips]
        else:
            ips = ips

        aliases = props.get('aliases', [])
        comment = props.get('comment', None)
        public = props.get('public', None)
        hs.append(Host(canonical, ips, aliases, public, comment))
    return hs


def patch_hosts(hosts, patchfile):
    """Patch a list of hosts in place"""
    patch = yaml.load(open(patchfile))
    
    hs = {}
    for h in hosts:
        hs[h.canonical] = h

    for k, changes in patch.items():
        h = hs[k]
        for change in changes:
            mo = re.match(r'(?P<verb>Add|Remove) alias (?P<subject>.*)', change)
            if not mo:
                raise ValueError("Error parsing patch file: %s" % change)
            if mo.group('verb') == 'Add':
                h.aliases.append(mo.group('subject'))
            else:
                h.aliases.remove(mo.group('subject'))
                

from iptools import *

class GlobalRoute(object):
    """Represents all routeable IP addresses"""
    def __init__(self):
        self.non_routeable = IpRangeList(
            '127/8',
            '10/8',
            '192.168/16',
            '172.16/12',
        )

    def __contains__(self, ip):
        return ip not in self.non_routeable


def hosts_to_zonefile(hosts, zone_name, routes=[GlobalRoute()], admin=None, soa=None, only_routeable=False, template=None):
    zone = Zone(zone_name, admin=admin, soa=soa)

    if template is None:
        template = os.path.join('templates/', zone_name)
    try:
        zone.add_from_template(template)
    except IOError:
        print >>sys.stderr, "Warning: Couldn't open zone template", template

    for host in hosts:
        try:
            ip = host.routeable_ip(routes)
        except NoRouteToHost:
            if only_routeable:
                continue
            # If there is no route, use the first IP in the list
            # This will ensure we get a "No route to host" error rather than an "unknown host" error.
            try:
                ip = host.ips[0]
            except IndexError:
                continue

        if host.canonical in zone:
            zone.add_a(host.canonical, ip)
        for alias in host.aliases:
            if alias in zone:
                zone.add_cname(alias, host.canonical)

    return zone


def reverse_zone(hosts, zone_name, admin, soa=None, template=None):
    zone = ReverseZone(zone_name, admin=admin, soa=soa)

    if template is None:
        template = os.path.join('templates/', zone_name)
    try:
        zone.add_from_template(template)
    except IOError:
        print >>sys.stderr, "Warning: Couldn't open zone template", template

    for host in hosts:
        for rip in host.reverse_ips():
            if rip in zone:
                zone.add_ptr(rip, host.canonical)
    return zone


def networks(networkfile='networks.yaml'):
    nets = yaml.load(open(networkfile))
    network_routes = {}
    for k, rs in nets.items():
        routes = []
        for r in rs:
            if r == '*':
                routes.append(GlobalRoute())
            else:
                routes.append(IpRange(r))
        network_routes[k] = routes
    return network_routes


def public_view(hosts):
    """Convert a list of hosts to the set of those hosts that are publicly visible."""
    pub = []
    r = GlobalRoute()
    for host in hosts:
        if host.public:
            pub.append(Host(host.public, [ip for ip in host.ips if ip in r], host.public))
    return pub 


if __name__ == '__main__':
    from optparse import OptionParser

    parser = OptionParser(usage="%prog [OPTIONS...] [-a <admin>] -n <network> <zone>")
    parser.add_option('-n', '--network', help="Network (one of those defined in networks.yaml) [default: global internet]", default='*')
    parser.add_option('-c', '--host-config', help="Host configuration YAML file to use [default: hosts.yaml]", default="hosts.yaml")
    parser.add_option('-a', '--admin', help="Admin e-mail address for the zone [default: hostmaster@<zone>]")
    parser.add_option('-s', '--soa', help="Authoritative nameserver for the zone [default: first nameserver in template]")
    parser.add_option('-p', '--patch', help="Pre-apply a patch file to the host configuration")
    parser.add_option('-P', '--public', help="Use public names", action='store_true')
    parser.add_option('-r', '--only-routeable', help="Only include routeable machines in the zone (ignored in reverse zones)", action="store_true")
    parser.add_option('-S', '--serial', help="Serial number; if unspecified, output a template for later substitution")
    parser.add_option('-t', '--template', help="Use a specific template file as the basis for a zone")

    options, args = parser.parse_args()

    if len(args) == 0:
        parser.error("No zone specified.")
    if len(args) > 1:
        parser.error("Wrong number of command line arguments.")

    zone = args[0]
    hosts = load_hosts(options.host_config)

    if options.patch:
        patch_hosts(hosts, options.patch)

    if options.public:
        hosts = public_view(hosts)

    if zone.endswith('.in-addr.arpa'):
        if options.admin is None:
            parser.error("You must provide an admin e-mail address (-a) when defining a reverse DNS zone.")

        zone = reverse_zone(hosts, zone, admin=options.admin, soa=options.soa, template=options.template)
    else:
        network_routes = networks()

        try:
            routes = network_routes[options.network]
        except KeyError:
            routes = [GlobalRoute()]

        zone = hosts_to_zonefile(hosts, zone, routes=routes, admin=options.admin, soa=options.soa, only_routeable=options.only_routeable, template=options.template)
        zone.check()
    if options.serial:
        if not options.serial.isdigit() or not (0 <= int(options.serial) < 2 ** 32):
            parser.error("Serial number must be a 32-bit integer.")
        zone.write_zonefile(sys.stdout, options.serial)
    else:
        zone.write_template(sys.stdout)

